#!/usr/bin/env python3
# Copyright (c) 2025 Huawei Device Co., Ltd.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


import argparse
import os
import subprocess
import sys
from collections.abc import Sequence
from pathlib import Path
from subprocess import CalledProcessError


class Color:
    """ANSI escape codes used to colorize log output in the terminal.

    RST  : Reset all attributes.
    INV  : Invert foreground and background colors.
    RED  : Red foreground (used for errors).
    GREEN: Green foreground (used for success or stages).
    BLUE : Blue foreground (used for actions).
    GREY : Grey foreground (used for verbose/debug logs).
    """

    RST = "\033[0m"
    INV = "\033[7m"
    RED = "\033[31m"
    GREEN = "\033[32m"
    BLUE = "\033[34m"
    GREY = "\033[90m"


def log_error(msg: str):
    print(f"{Color.INV}{Color.RED}[!]{Color.RST}", msg, file=sys.stderr)


def log_stage(msg: str):
    print(f"{Color.INV}{Color.GREEN}[*]{Color.RST}", msg, file=sys.stderr)


def log_action(msg: str):
    print(f"{Color.INV}{Color.BLUE}[-]{Color.RST}", msg, file=sys.stderr)


def log_verbose(msg: str):
    print(
        f"{Color.INV}{Color.GREY}[.]{Color.RST}",
        f"{Color.GREY}{msg}{Color.RST}",
        file=sys.stderr,
    )


def _log_command(command: Sequence[Path | str], cwd: Path | str):
    log_action(" ".join(str(c) for c in command))
    log_verbose(f"working directory: {cwd}")


def run_command(
    command: Sequence[Path | str],
    cwd: Path | str,
    check: bool = True,
):
    _log_command(command, cwd)
    subprocess.run(command, cwd=cwd, check=check)


def run_command_with_output(
    command: Sequence[Path | str],
    cwd: Path | str,
    check: bool = True,
) -> str:
    _log_command(command, cwd)
    result = subprocess.run(
        command,
        cwd=cwd,
        check=check,
        capture_output=True,
        text=True,
    )
    return result.stdout.strip()


def get_project_root() -> Path:
    return Path(__file__).parent.parent.resolve()


class PythonEnv:
    """Context manager for activating and restoring a Python virtual environment."""

    def __init__(self, project_root: Path):
        self.project_root = project_root
        self.compiler_path = project_root / "compiler"
        self.venv_bin = project_root / ".venv/bin"
        self._original_env: dict[str, str] = {}

    def __enter__(self):
        if not self.venv_bin.exists():
            log_error(f"virtual environment directory not found: {self.venv_bin}")
            sys.exit(1)

        # Save current environment
        self._original_env = {
            "PATH": os.environ.get("PATH", ""),
            "PYTHONPATH": os.environ.get("PYTHONPATH", ""),
        }

        # Modify environment
        os.environ["PYTHONPATH"] = str(self.compiler_path)
        os.environ["PATH"] = str(self.venv_bin) + os.pathsep + os.environ["PATH"]
        log_verbose(f"Python env activated: {self.venv_bin}")

        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """Restore original environment."""
        if not self._original_env:
            log_verbose("exit_py_env called before activation, skipping.")
            return

        os.environ["PATH"] = self._original_env["PATH"]
        if self._original_env["PYTHONPATH"]:
            os.environ["PYTHONPATH"] = self._original_env["PYTHONPATH"]
        else:
            os.environ.pop("PYTHONPATH", None)

        log_verbose("Python env deactivated.")


def run_module(module_name: str, args: Sequence[Path | str], project_root: Path):
    coverage_run = [
        "coverage",
        "run",
        "--parallel-mode",
        "--module",
        module_name,
        *args,
    ]
    run_command(coverage_run, cwd=project_root)


def run_taihe_tryit(args: Sequence[Path | str], project_root: Path):
    run_module("taihe.cli.tryit", args, project_root)


def run_pytest(project_root: Path, **config: Path | str):
    log_stage("Run PyTest...")
    run_module("pytest", [project_root / "compiler/tests"], project_root)


def run_core_tests(project_root: Path, **config: Path | str):
    log_stage("Run core tests...")
    for test_dir in [project_root / "test/rgb", project_root / "test/object"]:
        log_stage(f"Testing: {test_dir}")
        run_taihe_tryit(["test", "-u", "cpp", test_dir], project_root)

def run_ani_test_build(
    project_root: Path,
    build_dir_name: str,
    enable_coverage: bool = True,
    taihec_path: Path | None = None,
    panda_home: Path | str | None = None
):
    build_dir = project_root / build_dir_name

    if not build_dir.exists():
        log_verbose("Build directory not found. Running CMake configuration.")
        ninja_cmd = ["cmake", "-B", str(build_dir), "-GNinja"]
        if enable_coverage:
            ninja_cmd.append("-DENABLE_COVERAGE=ON")
        if taihec_path:
            ninja_cmd.append(f"-DBIN_TAIHEC_PATH={taihec_path}")
        if panda_home:
            ninja_cmd.append(f"-DPANDA_HOME={panda_home}")
        run_command(ninja_cmd, cwd=project_root)
    else:
        log_verbose("Build directory exists. Skipping configuration.")

    build_cmd = ["cmake", "--build", str(build_dir), "--verbose"]
    run_command(build_cmd, cwd=project_root)

def run_ani_tests(project_root: Path, **config: Path | str):
    log_stage("Run ANI tests...")
    run_ani_test_build(
        project_root,
        build_dir_name="build",
        enable_coverage=True,
        panda_home=config.get("panda_home")
    )

def run_bin_ani_tests(project_root: Path, **config: Path | str):
    log_stage("Run BIN-ANI tests...")

    taihe_bin_tar_name = "taihe-linux-x86_64-0.55.5-20251125.tar.gz"
    taihe_bin_tar_url = f"https://gitcode.com/tongdiaoZS/taihe-bin/releases/download/0.55.5/{taihe_bin_tar_name}"
    taihe_bin_tar_path = project_root / taihe_bin_tar_name
    taihe_bin_extract_dir = project_root / "taihe_bin"

    if not taihe_bin_tar_path.exists():
        log_verbose("Downloading Taihe binary package...")
        try:
            run_command(["curl", "-fSL", "-o", str(taihe_bin_tar_path), taihe_bin_tar_url], cwd=project_root)
        except CalledProcessError as e:
            log_error(f"Failed to download Taihe binary package from {taihe_bin_tar_url}")
            log_error(f"curl exited with code {e.returncode}")
            raise
    else:
        log_verbose("Taihe binary already exists. Skipping download.")

    if not taihe_bin_extract_dir.exists():
        log_verbose("Extracting Taihe binary package...")
        run_command(["mkdir", "-p", str(taihe_bin_extract_dir)], cwd=project_root)
        run_command(["tar", "-xzf", str(taihe_bin_tar_path), "-C", str(taihe_bin_extract_dir)], cwd=project_root)
    else:
        log_verbose("Taihe binary directory already exists. Skipping extraction.")

    taihe_bin_path = taihe_bin_extract_dir / "taihe" / "bin" / "taihec"

    if not taihe_bin_path.exists():
        raise FileNotFoundError(f"taihec not found at {taihe_bin_path}")

    run_ani_test_build(
        project_root,
        build_dir_name="build_bin",
        enable_coverage=False,
        taihec_path=taihe_bin_path,
        panda_home=config.get("panda_home")
    )


def run_cmake_test(project_root: Path, **config: Path | str):
    log_stage("Run CMake test...")

    test_dir = project_root / "test/cmake_test"
    run_taihe_tryit(["generate", "-u", "sts", test_dir, "-Bcmake"], project_root)

    try:
        taihec_cmd = ["taihec", "--print-panda-vm-path"]
        panda_vm_path = run_command_with_output(taihec_cmd, cwd=project_root)
        ani_path = Path(panda_vm_path) / "ohos_arm64/include/plugins/ets/runtime/ani"
    except (subprocess.CalledProcessError, FileNotFoundError) as e:
        log_error(f"Unable to get panda-vm-path: {e}")
        return

    ninja_cmd = [
        "cmake",
        "-B",
        "build/generated",
        f"-DEXTERNAL_INCLUDE={ani_path}",
    ]
    run_command(ninja_cmd, cwd=test_dir)
    build_cmd = ["cmake", "--build", "build/generated"]
    run_command(build_cmd, cwd=test_dir)


def generate_coverage_report(project_root: Path):
    log_stage("Generate coverage report...")
    try:
        run_command(["coverage", "combine"], cwd=project_root)
        run_command(["coverage", "report"], cwd=project_root)
    except subprocess.CalledProcessError as e:
        log_error(f"Failed to generate coverage report: {e}")


TESTS = {
    "pytest": run_pytest,
    "core": run_core_tests,
    "ani": run_ani_tests,
    "bin_ani": run_bin_ani_tests,
    "cmake": run_cmake_test,
}

MODE_TO_TESTS = {
    "github-push-ci": ["pytest", "core", "ani", "cmake"],
    "github-pr-ci": ["pytest", "core", "ani", "cmake"],
    "ggw-ci": ["pytest", "ani", "bin_ani"],
}


def main():
    parser = argparse.ArgumentParser(
        description="Taihe integration test",
        formatter_class=argparse.RawTextHelpFormatter,
    )
    parser.add_argument("--panda-home", type=str, help="Specify PANDA_HOME path")
    group = parser.add_mutually_exclusive_group()
    group.add_argument(
        "--run",
        nargs="*",
        choices=TESTS.keys(),
        help="Specify which test(s) to run",
    )
    modes = ", ".join(MODE_TO_TESTS.keys())
    group.add_argument(
        "--mode",
        type=str,
        choices=MODE_TO_TESTS.keys(),
        help=f"Run tests under mode ({modes})",
    )
    args = parser.parse_args()

    project_root = get_project_root()

    if args.mode:
        selected_tests = MODE_TO_TESTS[args.mode]
    elif args.run:
        selected_tests = args.run
    else:
        selected_tests = ["pytest", "core", "ani", "cmake"]

    only_bin_ani = all(name == "bin_ani" for name in selected_tests)

    for name in selected_tests:
        log_stage(f"=== Running test: {name} ===")

        if name == "bin_ani":
            TESTS[name](project_root, panda_home=args.panda_home)
        else:
            with PythonEnv(project_root):
                TESTS[name](project_root, panda_home=args.panda_home)

    if not only_bin_ani:
        with PythonEnv(project_root):
            generate_coverage_report(project_root)
    else:
        log_verbose("Only 'bin_ani' tests run, skipping coverage report.")


if __name__ == "__main__":
    main()
